#!/usr/bin/env python2.7

"""dx-docker Run Docker images in your DNAnexus app."""

import argparse
import sys
import os
import errno
import subprocess
import pprint
import json
import multiprocessing.dummy
import tempfile
import glob
import datetime
import time
import dateutil.parser
import dxpy.utils.printing
import dxpy
import urllib
import tarfile
import shutil

CACHE_DIR = '/tmp/dx-docker-cache'

def shell(cmd, ignore_error=False):
    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError as e:
        if ignore_error:
            return
        else:
            sys.exit(e.returncode)

def shell_suppress(cmd, ignore_error=False):
    out = ""
    try:
        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        if ignore_error:
            pass
        else:
            print e.output.strip()
            raise
    return out

def makedirs(path):
    try:
        os.makedirs(path)
    except OSError as e:  # If the directory already exists, continue gracefully
        if e.errno != errno.EEXIST:
            raise
        pass

parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers()

def get_aci_fname(image):
    matching_files = glob.glob("{cachedir}/{image}.aci".format(cachedir=CACHE_DIR, image=urllib.quote_plus(image)))
    if len(matching_files) == 0:
        return None
    elif len(matching_files) == 1:
        return matching_files[0]
    else:
        raise Exception("dx-docker cache error: Exactly one aci archive should correspond to an image name: {} found archives".format(len(matching_files)))

def extract_aci(acifile):
    tmpdir = tempfile.mkdtemp()
    shell_suppress(["tar", "-xf", acifile, "--exclude", "rootfs/dev/*", "-C", tmpdir], ignore_error=True)
    return tmpdir

def retry(tries=3, delay=1, backoff=2, retry_on_exceptions=None):
    '''A decorator that retries a function until it completes without throwing an exception.'''

    if tries < 0:
        raise ValueError("tries must be a non-negative number")
    if delay <= 0:
        raise ValueError("delay must be a positive number")
    if backoff <= 1:
        raise ValueError("backoff must be greater than 1")
    if not retry_on_exceptions:
        retry_on_exceptions = (Exception,)

    def decorated_f(f):
        def wrapper(*args, **kwargs):
            mdelay = delay
            for i in range(tries):
                try:
                    f(*args, **kwargs)
                except retry_on_exceptions as retry_exception:
                    sys.stderr.write(dxpy.utils.printing.YELLOW("Error. Retrying in {} s.\n".format(mdelay)))
                    time.sleep(mdelay)
                    mdelay *= backoff
                    if i == tries - 1:
                        raise retry_exception
                    continue
                break
        return wrapper
    return decorated_f

def get_aci(image, imgname, prefix="", quiet=False):
    def docker2aci_line(cmd):
        popen = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=open(os.devnull, "w"), universal_newlines=True)
        stderr_lines = iter(popen.stderr.readline, "")
        for stderr_line in stderr_lines:
            yield stderr_line

        popen.stderr.close()
        return_code = popen.wait()
        if return_code != 0:
            raise subprocess.CalledProcessError(return_code, cmd)
    tmpdir = tempfile.mkdtemp()
    os.chdir(tmpdir)
    cmd = ["docker2aci", "-compression", "none", imgname]
    try:
        if quiet:
            shell_suppress(cmd)
        else:
            layers = {}
            for line in docker2aci_line(cmd):
                fields = line.rstrip().split()
                if len(fields) == 7 and fields[0] == "Downloading":
                    _, imghash, downloaded, unita, _, total, unitb = line.rstrip().split()
                else:
                    sys.stderr.write(line)
                    continue
                percent = round(float(downloaded)/float(total),1)
                def print_progress():
                    sys.stderr.write(dxpy.utils.printing.BLUE(imghash)+": "+dxpy.utils.printing.YELLOW("{} {}/{} {}\n".format(downloaded, unita, total, unitb)))
                if imghash not in layers:
                    layers[imghash] = percent
                    if percent == 1:
                        print_progress()
                elif layers[imghash] != percent and percent > 0:
                    print_progress()
                    layers[imghash] = percent
    except Exception as e:
        raise Exception("Failed to obtain image {}.  Does this image exist in the specified repository?\n\n{}".format(image, e))

    img_dirname = prefix+CACHE_DIR
    makedirs(img_dirname)
    os.rename(glob.glob(tmpdir+"/*.aci")[0], img_dirname+"/"+urllib.quote_plus(image)+".aci")
    shutil.rmtree(tmpdir)

def register_docker_subcommand(cmd):
    orig_user_agent = dxpy.USER_AGENT
    dxpy.USER_AGENT += " dx-docker {}".format(cmd)
    dxpy.api.system_greet()
    dxpy.USER_AGENT = orig_user_agent

parser_pull = subparsers.add_parser('pull', help="Pulls a docker image for use in DNAnexus")
parser_pull.add_argument("image", help="image name")
parser_pull.add_argument("-q", "--quiet", action='store_true', help="Suppress printing pull progress to stderr")
@retry()
def pull(args):
    if not os.path.isdir(CACHE_DIR):
        makedirs(CACHE_DIR)

    get_aci(args.image, "docker://"+args.image, quiet=args.quiet)
parser_pull.set_defaults(func=pull)

parser_run = subparsers.add_parser('run', help="Runs a docker image in a container")
parser_run.add_argument("-v", "--volume", help="Directory to mount inside the container. Can be supplied multiple times (e.g. -v /mnt/data:/mnt/data -v /host:/guest", action='append', default=[])
parser_run.add_argument("-e", "--env", help="Environment variables to set within container. Can be supplied multiple times (e.g. -e foo=bar -e pizza=pie", action='append', default=[])
parser_run.add_argument("--entrypoint",  help="Overwrite default entry point for image")
parser_run.add_argument("-w", "--workdir",  help="Working directory")
parser_run.add_argument("-q", "--quiet",  action='store_true', help="Suppress printing of image metadata")
parser_run.add_argument("--rootfs",  help="Use directory pointed to here for rootfs instead of extracting the image (for expert use/development purposes)")
parser_run.add_argument("--rm",  action="store_true", help="Automatically remove the container when it exits")
parser_run.add_argument("image", help="image name")
parser_run.add_argument("command", help="command to run within container", nargs=argparse.REMAINDER, default=[])
def run(args):
    register_docker_subcommand("run")
    acifile = get_aci_fname(args.image)
    if not acifile:
        pull(parser_pull.parse_args([args.image]))
        acifile = get_aci_fname(args.image)

    if args.rootfs:
        tmpdir = args.rootfs
    else:
        tmpdir = extract_aci(acifile)

    with open("{tmpdir}/manifest".format(tmpdir=tmpdir)) as f:
        imagemeta = json.loads(f.read())

    annotations = {a['name'].split("/")[-1]: a['value'] for a in imagemeta['annotations']}
    labels = {a['name'].split("/")[-1]: a['value'] for a in imagemeta['labels']}
    if not args.quiet:
        def print_annotation(name, key, data=annotations):
            if key in data:
                sys.stderr.write(dxpy.utils.printing.BLUE(name)+"\t"+dxpy.utils.printing.YELLOW(data[key])+"\n")

        stats = {
            'size': shell_suppress(["du", "-sh", tmpdir]).split()[0]
        }

        if 'created' in annotations:
            created_time = dateutil.parser.parse(annotations['created'])
            current_time = datetime.datetime.now(dateutil.tz.tzutc())
            time_since_creation = current_time - created_time
            if (time_since_creation.days == 0 and time_since_creation.seconds < 10*60):
               human_readable_delta = "(just now)"
            else:
               human_readable_delta = "({}d {}h ago)".format(time_since_creation.days, time_since_creation.seconds/3600)
            stats['created'] = annotations['created'] + ' ' + human_readable_delta

        print_annotation("Image Repo", 'repository')
        print_annotation("Image Tag", 'version', data=labels)
        print_annotation("Image Size", 'size', data=stats)
        print_annotation("Image ID", 'imageid')
        print_annotation("Parent ID", 'parentimageid')
        print_annotation("Last Updated", 'created', data=stats)
        print_annotation("Registry", 'registryurl')
        print_annotation("ENTRYPOINT", 'entrypoint')
        if 'app' in imagemeta:
            print_annotation("Working Dir", 'workingDirectory', data=imagemeta['app'])
        print_annotation("Default CMD", 'cmd')
        sys.stderr.write("\n\n")

    vols = ["/run/shm/:/dev/shm"] + args.volume
    volume = []
    for v in vols:
       volume.append("-b")
       volume.append(v)

    entrypoint = []
    if args.entrypoint:
        entrypoint = [args.entrypoint]
    elif 'entrypoint' in annotations:
        entrypoint = json.loads(annotations['entrypoint'])

    cmd = []
    if args.command:
         cmd = args.command
    elif 'cmd' in annotations and not args.entrypoint:
         cmd = json.loads(annotations['cmd'])

    container_cmd = entrypoint + cmd

    workdir = "/"
    if 'app' in imagemeta and 'workingDirectory' in imagemeta['app']:
       workdir = imagemeta['app']['workingDirectory']
    if args.workdir:
       workdir = args.workdir

    default_homedir_env = "HOME=/root"
    env = ["env", "-i", default_homedir_env]
    if 'app' in imagemeta and 'environment' in imagemeta['app']:
        env.extend([var['name'] +"="+var['value']  for var in imagemeta['app']['environment']] + args.env)

    if args.rm and not args.quiet:
        sys.stderr.write(dxpy.utils.printing.YELLOW("--rm: Note that by default the container is not preserved after execution\n"))

    # A list of dirs and files of the host rootfs that will be accessible in the confined
    # environment just as if they were part of the guest rootfs. This list is based on the
    # proot -R option: https://github.com/proot-me/PRoot/blob/master/doc/proot/manual.txt
    to_mount = [ '/etc/host.conf', '/etc/hosts', '/etc/mtab', '/etc/networks', '/etc/passwd',
                 '/etc/group', '/etc/nsswitch.conf', '/etc/resolv.conf', '/etc/localtime',
                 '/dev/', '/sys/', '/proc/', '/run/', '/var/run/dbus/system_bus_socket']
    bounded_dirs = []
    for b in to_mount:
        bounded_dirs.append('-b')
        bounded_dirs.append(b)

    # TODO: implement some form of nested subprocesses to avoid quoting command array if possible
    proot_cmd =  ["proot", '-r', '{tmpdir}/rootfs'.format(tmpdir=tmpdir)] + \
                 bounded_dirs + \
                 volume + \
                 ["-w", workdir] + \
                 env + ["/bin/sh", "-c"] + \
                 [" ".join(("'"+x+"'" for x in container_cmd))]

    shell(proot_cmd)
    if not args.rootfs:
        shutil.rmtree(tmpdir)
parser_run.set_defaults(func=run)

def local_image_to_aci(image, tmpdir, dirprefix, alternative_export=False):
    # Ensure core utilities are available on local host
    shell_suppress(["which", "docker"])
    shell_suppress(["which", "docker2aci"])

    # Check if image is docker image id or digest
    def is_image_id_or_digest(image):
        # Remove prefix for string comparisons
        if ('@sha256:' in image or image.startswith('sha256:')):
            img_id = image.split('sha256:')[1]
        else:
            img_id = image

        image_id_list = shell_suppress(["docker", "images", "-q", "--no-trunc"]).split("\n")
        digest_list = shell_suppress(["docker", "images", "--digests", "--filter", "dangling=false", "--format", "{{.Repository}}@{{.Digest}}"]).split("\n")

        # Check if given image id equal to full hash or short hash (first 12 chars)
        if any(img_id in [image_id.split('sha256:')[1], image_id.split('sha256:')[1][:12]] for image_id in image_id_list if image_id.strip()):
            sys.stderr.write(dxpy.utils.printing.YELLOW("Image ID: {} found in local docker cache\n".format(image)))
            return True
        # Check if given image id is a digest
        elif any(image == digest for digest in digest_list if digest.strip()):
            sys.stderr.write(dxpy.utils.printing.YELLOW("Digest: {} found in local docker cache\n".format(image)))
            return True
        else:
            return False

    # We have observed an instance where the Docker CLI and docker2aci do not properly export a rootfs.  Here is a hopefully temporary workaround
    if alternative_export:
        # TODO: Do obvious things to make the following steps more efficient (e.g. no re-untarring rootfs, ...)
        sys.stderr.write(dxpy.utils.printing.YELLOW("Running alternative export for {}\n".format(image)))
        shell_suppress(["docker", "run", "--entrypoint", "tar", "-w", "/tmp", "-v", "{}:/tmp".format(tmpdir), image, "-cf", "rootfs.tar", "/"], ignore_error=True)
        sys.stderr.write(dxpy.utils.printing.YELLOW("Extracting root file system for {}\n".format(image)))
        shell(["mkdir", "-p", tmpdir+"/new-aci-template/rootfs"])
        shell_suppress(["tar", "-xf", tmpdir+"/rootfs.tar", "--exclude", "dev/*", "--exclude", "proc/*", "--exclude", "sys/*", "-C", tmpdir+"/new-aci-template/rootfs/"], ignore_error=True)
        shell(["rm", "-f", tmpdir+"/rootfs.tar"])
        # We do not have a manifest file so we must create one on our own from 'docker inspect'
        inspect_out = json.loads(shell_suppress(["docker", "inspect", image]))[0]
        imsplit = inspect_out["RepoTags"][0].split(':')
        dictifiedenv = [{"name": var.split("=")[0], "value": var.split("=")[1]} for var in inspect_out["Config"]["Env"]]
        entrypoint = inspect_out["Config"]["Entrypoint"]
        if not entrypoint:
            entrypoint = []
        workdir = inspect_out["Config"]["WorkingDir"]
        if not workdir or workdir == "":
            workdir = "/"

        # TODO if this needs to be kept longer-term: resolve simultaneous CMD + ENTRYPOINT and whether it goes in exec field (ignoring user/group ...)
        manifest={"acKind":"ImageManifest",
                  "acVersion":"0.8.5",
                  "name":imsplit[0],
                  "labels":[{"name":"version","value":imsplit[1]},
                            {"name":"os","value":inspect_out["Os"]},
                            {"name":"arch","value":inspect_out["Architecture"]}],
                             "app":{"exec":entrypoint,"user":"0","group":"0","workingDirectory":workdir,"environment":dictifiedenv},
                             "annotations":[{"name":"created","value":inspect_out["Created"]},
                                            {"name":"appc.io/docker/repository","value":imsplit[0]},
                                            {"name":"appc.io/docker/imageid","value":inspect_out["Id"]},
                                            {"name":"appc.io/docker/parentimageid","value":inspect_out["Parent"]},
                                            {"name":"appc.io/docker/entrypoint","value":json.dumps(entrypoint)},
                                            {"name":"appc.io/docker/tag","value":imsplit[1]}]}
        with open(tmpdir+"/new-aci-template/manifest", "w") as f:
            f.write(json.dumps(manifest))
        shell(["mkdir", "-p", dirprefix+CACHE_DIR])
        sys.stderr.write(dxpy.utils.printing.YELLOW("Creating ACI filesystem bundle for {}\n".format(image)))
        shell_suppress(["tar", "-C", tmpdir+"/new-aci-template", "-cf", dirprefix+CACHE_DIR+"/"+urllib.quote_plus(image)+".aci", "."], ignore_error=True)
    # Use Docker and docker2aci to generate an aci
    else:
        sys.stderr.write(dxpy.utils.printing.YELLOW("Exporting Docker image {}\n".format(image)))
        docker_tarball = tmpdir+"/image.docker.tgz"
        # If image id is given instead of 'repo:tag', tag image as 'dx-docker:id' and export with `docker save dx-docker:id`
        if is_image_id_or_digest(image):
            # Remove 'sha256:' prefix if it is supplied, colon cannot exist in the repo:tag format
            if ('@sha256:' in image or image.startswith('sha256:')):
                image_id_tag = image.split('sha256:')[1]
            else:
                image_id_tag = image
            shell_suppress(["docker", "tag", image, "dx-docker:{}".format(image_id_tag)])
            shell(["docker", "save", "-o", docker_tarball, "dx-docker:{}".format(image_id_tag)])
            shell_suppress(["docker", "rmi", "dx-docker:{}".format(image_id_tag)])
        else:
            shell(["docker", "save", "-o", docker_tarball, image])
        get_aci(image, docker_tarball, prefix=dirprefix)

parser_add_to_applet = subparsers.add_parser('add-to-applet', help="Adds a local Docker image to an applet")
parser_add_to_applet.add_argument("image", help="image name")
parser_add_to_applet.add_argument("applet", help="directory corresponding to applet")
parser_add_to_applet.add_argument("--alternative_export", help="EXPERT ONLY: Use alternative method to export Docker image since Docker CLI export sometimes doesn't create the root filesystem properly.", action="store_true")
def add_to_applet(args):
    register_docker_subcommand("add-to-applet")
    try:
        with open(args.applet+"/dxapp.json") as f:
            json.loads(f.read())
    except Exception as e:
        raise Exception("{} does not appear to have a dxapp.json that parses. Please make sure you have selected a working DNAnexus applet directory\n\n{}.".format(args.applet, e))
    tmpdir = tempfile.mkdtemp()
    local_image_to_aci(args.image, tmpdir, os.path.abspath(args.applet)+"/resources", args.alternative_export)
    shutil.rmtree(tmpdir)
parser_add_to_applet.set_defaults(func=add_to_applet)

def _upload_file(filename, project_id=None, folder_path=None, wait_on_close=True, hidden=True):
    kwargs = {'project': project_id, 'folder': folder_path}
    try:
        return dxpy.upload_local_file(filename, wait_on_close=wait_on_close, hidden=hidden, **kwargs).get_id()
    except Exception as e:
        print "Unable to upload file {}.\n{}.".format(filename, e)
        sys.exit(1)


parser_create_asset = subparsers.add_parser('create-asset', help="Caches a local Docker image as an asset in the DNAnexus platform (subject to change)")
parser_create_asset.add_argument("--output_path", "-o", help="Project ID and path in project to upload image to (defaults to project root)")
parser_create_asset.add_argument("image", help="image name")
parser_create_asset.add_argument("--alternative_export", help="EXPERT ONLY: Use alternative method to export Docker image since Docker CLI export sometimes doesn't create the root filesystem properly.", action="store_true")
def create_asset(args):
    register_docker_subcommand("create-asset")
    # Create asset
    tmpdir = tempfile.mkdtemp()
    asset_path = tmpdir+"/resources/usr/bin"
    makedirs(asset_path)

    local_image_to_aci(args.image, tmpdir, tmpdir+"/resources", args.alternative_export)

    imagename = args.image.replace("/", "%").replace(":","#")
    # Create the asset configuration
    dxasset = { 'name': imagename }
    dxasset['title'] = "DNAnexus asset for Docker image "+args.image
    dxasset['description'] = dxasset['title']
    dxasset['version'] = "0.0.1"
    dxasset['distribution'] = "Ubuntu"
    dxasset['release'] = "14.04"
    dxasset['instanceType'] = "mem2_ssd1_x4"
    with open("{}/dxasset.json".format(tmpdir), "w") as f:
        f.write(json.dumps(dxasset))

    project_id = ""
    folder_path = ""
    if args.output_path:
        if ":" in args.output_path:
            project_id, folder_path = args.output_path.split(":")
        else:
            project_id = dxpy.PROJECT_CONTEXT_ID
            folder_path = args.output_path
    else:
        project_id = dxpy.PROJECT_CONTEXT_ID
        folder_path = ""
    if not folder_path.startswith("/"):
        folder_path = "/" + folder_path

    sys.stderr.write(dxpy.utils.printing.YELLOW("Building DNAnexus asset for {}\n".format(args.image)))
    escaped_image_name = args.image.replace("/","\\/").replace(":", "\\:")
    image_tarball = imagename+".tar.gz"
    os.chdir(tmpdir)
    with tarfile.open(image_tarball, "w:gz") as tar:
        tar.add(tmpdir+"/resources", "/")
    sys.stderr.write(dxpy.utils.printing.YELLOW("Uploading DNAnexus asset for {}\n".format(args.image)))
    asset_tarball_id = _upload_file(image_tarball, project_id, folder_path)
    record_name = args.image
    record_details = {"archiveFileId": {"$dnanexus_link": asset_tarball_id}}
    record_properties = {
                          "title": dxasset["title"],
                          "description": dxasset["description"],
                          "version": "0.0.1",
                          "distribution": dxasset["distribution"],
                          "release": dxasset["release"]
                        }
    asset_bundle = dxpy.new_dxrecord(name=record_name,
                                     project=project_id, folder=folder_path,
                                     types=["AssetBundle"], details=record_details,
                                     properties=record_properties, close=True)

    # Add a property called {"AssetBundle": record-xxx} to the hidden tarball
    asset_file = dxpy.DXFile(asset_tarball_id, project=project_id)
    asset_file.set_properties({"AssetBundle": asset_bundle.get_id()})

    # Print output message for successful building of asset
    sys.stderr.write(dxpy.utils.printing.YELLOW("Image {} successfully cached in DNAnexus platform.\n".format(args.image)))
    sys.stderr.write("To include this cached image in an application, please include the following within the runspec/assetDepends list in your dxapp.json.\n")
    sys.stderr.write("    {"+dxpy.utils.printing.BLUE("""
        "project": \"{}\",
        "folder": \"{}\",
        "name": \"{}\",
        "version": "0.0.1"
    """).format(project_id, folder_path, args.image)+"}\n")

    # Remove the temporary directory
    shutil.rmtree(tmpdir)

parser_create_asset.set_defaults(func=create_asset)


if __name__ == "__main__":
    args = parser.parse_args()
    args.func(args)
